#include "./velocity.h"
#include "../constants.h"
#include "../sys/rand.h"
#include <math.h>

Real mds_velocity_kinetic_energy( const mds_vector_set velocity, size_t n )
{
	assert(velocity!=NULL);
	assert(n>0);
	{
		Real              K = 0;
		const mds_vector *v = velocity;
		while( n-- > 0 )
		{
			K += v->info.mass * mds_vector_normsq( v );
			++v;
		}
		
		return MDS_F(0.5)*K;
	}
}

Real mds_velocity_initialize( const mds_vector_set position, mds_vector_set velocity, size_t n, Real Temperature )
{
	const Real RT    = MDS_R * Temperature;
	const size_t dof =  3*n; 
	const Real K0    =  dof * RT * MDS_F(0.5);
	assert(velocity!=NULL);
	assert(Temperature>=0);
	assert(n>0);
	fprintf( stderr, "K0=%g\n", K0);
	if( K0 > 0 )
	{
		Real K1  = 0;
		Real lam = 0;
		
		/*************************************************************/
		/*  generate a zero momentum Maxwell Boltzmann distribution */
		/***********************************************************/
		mds_vector P  = { 0, 0, 0 };
		mds_vector dP = { 0, 0, 0 }; 
		register size_t i;
		for( i=0; i < n; ++i )
		{
			const mds_vector *r   = &position[i];
			mds_vector       *v   = &velocity[i];
			/*-------------------------------------*/
			/* Get the Maxwell-Boltzmann variance */
			/*-----------------------------------*/
			const Real        sig = (Real)sqrt( RT * r->info.inv_mass );
			
			/*--------------------------*/
			/* Get a gaussian velocity */
			/*------------------------*/
			v->x = sig * mds_normal();
			v->y = sig * mds_normal();
			v->z = sig * mds_normal();
			
			/*-------------------------------*/			
			/* update total linear momentum */
			/*-----------------------------*/
			mds_vector_muladd( &P, v->info.mass, v);
		}
		mds_vector_mulset( &dP, -MDS_F(1.0) / n, &P );
		
		/*************************************************************/
		/*  dispatch dP and compute current kinetic energy          */
		/***********************************************************/
		for( i=0; i < n; ++i )
		{
			const mds_vector *r   = &position[i];
			mds_vector       *v   = &velocity[i];
			mds_vector_muladd( v, r->info.inv_mass, &dP );
			K1 +=  v->info.mass * mds_vector_normsq( v );
		}
		K1 *= MDS_F(0.5);
		
		/*************************************************************/
		/*  Rescale velocities                                      */
		/***********************************************************/
		lam = (Real)sqrt(K0/K1);
		K1  = 0;
		for( i=0; i < n; ++i )
		{
			mds_vector       *v   = &velocity[i];
			mds_vector_mul(v, lam );
			K1 +=  v->info.mass * mds_vector_normsq( v );
		}
		K1 *= MDS_F(0.5);
		
		fprintf( stderr, "K=%g / T=%g\n", K1, (2*K1)/(dof*MDS_R) );
		return K1;
	}
	else {
		size_t i;
		for(  i=0; i < n; ++i )
		{
			mds_vector_ldz( &velocity[i] );
		}
		
		return 0;
	}
	
	
}

Real mds_velocity_rescale( const mds_vector_set velocity, size_t n, const Real alpha )
{
	Real            K = 0;
	register size_t i;
	assert( velocity );
	assert( n > 0 );
	for( i=0; i < n; ++i )
	{
		mds_vector  *v = &velocity[i];
		mds_vector_mul( v , alpha );
		K += v->info.mass * mds_vector_normsq( v );
	}
	K *= MDS_F(0.5);
	return K;
}


